feat: mx9: add fused MX9 quantization kernel with pack/unpack#33
Open
Jiarwang77 wants to merge 2 commits into
Open
feat: mx9: add fused MX9 quantization kernel with pack/unpack#33Jiarwang77 wants to merge 2 commits into
Jiarwang77 wants to merge 2 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds real packed MX9 quantization support via a fused Triton pack/unpack kernel, enabling a compact 3-part representation (q / max_exp / prime) intended for inference/storage compression, along with an extensive unit test suite validating bit-exact round-trips against a clamp-127 reference.
Changes:
- Introduces
convert_to_mx9/convert_from_mx9Triton implementation for real packed MX9 (int8 q + E8M0-like exponent byte + prime bitmap). - Adds comprehensive MX9 pack/unpack tests (format, intermediates, round-trip, edge cases, invalid inputs).
- Adds
alto.kernels.mxpackage initializer and documentation.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
alto/kernels/mx/mx9_quantization.py |
New Triton MX9 pack/unpack kernel + Python wrappers for producing/consuming the packed 3-tensor representation. |
tests/unittest/mx9_mx6/test_mx9_quantization.py |
New end-to-end and component-level tests validating storage format and bit-exact round-trip behavior. |
alto/kernels/mx/__init__.py |
Declares the MX “real kernel” package and documents its relationship to the fake-quant reference path. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+435
to
+437
| assert out_dtype in _TORCH_TO_TL | ||
| assert block_size == 16, f"block_size only supports 16, got {block_size}" | ||
|
|
Comment on lines
+191
to
+197
| def test_max_exp_biased_range(): | ||
| # E8M0 stores true exponent + 127; finite fp32 true exponent in [-126, 127], | ||
| # biased to [1, 254]; should never be 0 or 255. | ||
| x = _rand((16, 64), torch.float32).cuda() | ||
| _, max_exp, _ = convert_to_mx9(x, block_size=16) | ||
| assert max_exp.min().item() >= 1 | ||
| assert max_exp.max().item() <= 254 |
hann-wang
reviewed
Jun 30, 2026
| # Copyright (c) 2026 Advanced Micro Devices, Inc. | ||
| # | ||
| # SPDX-License-Identifier: MIT | ||
| """MX real (fused Triton) kernels. |
Collaborator
There was a problem hiding this comment.
This is misleading. We don't have real mx gemm now.
|
|
||
| # Q part of QDQ operates on the *original* x (NaN will produce garbage, see below). | ||
| xf = x.to(tl.float32) | ||
| q = _round_half_even(xf / scale) |
Collaborator
There was a problem hiding this comment.
is it possible to use triton.language.div_rn?
Collaborator
Author
There was a problem hiding this comment.
Yes, I found "PyTorch eager uses the equivalent of Triton's div_rn" in torch inductor. It can be more accurate.
- patcher: format=="mx9" now dispatches to convert_to/from_mx9 (packed Triton kernel) instead of mx9_fake_quantize (pure PyTorch). Validated: LLaMA-3.2-1B wikitext loss 2.1141 (== fake-quant 2.1140). - mx9_quantization: use tl.div_rn for IEEE round-nearest division (fixes potential 2-ULP fast-div rounding at .5 boundaries); add dtype asserts on q/max_exp/prime in convert_from_mx9. - kernels/mx/__init__: update docstring.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add a fused Triton kernel for real packed MX9 quantization (
convert_to_mx9/convert_from_mx9), complementing the existing fake-quant path inalto.modifiers.quantization.mx.Unlike the fake-quant path (which outputs dequantized bf16/fp32 for training), this kernel produces a compact three-part packed representation:
q—[n_blocks, 16]int8: per-element quantized values, symmetric clamp ±127.max_exp—[n_blocks]uint8: per-block shared exponent in E8M0 format (true exponent + 127 bias).prime—[n_blocks, 1]uint8: per-block prime bitmap, 1 bit per pair of 2 elements, byte-packed.This achieves 9 bits/element (8b value + 0.5b amortized exponent + 0.5b prime), enabling real weight compression for inference/storage.
Key design choices
_t_exponent.Test coverage (70 tests, all passing)
Test plan
pytest tests/unittest/mx9_mx6/test_mx9_quantization.py— 70 passed (ROCm/HIP, MI300X)@triton_op+register_fakefortorch.compileintegration